import os
import pickle
import numpy as np
from tqdm import tqdm
from PIL import Image
from encoder.clip import clip_encoder

data_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'data'))
image_path = os.path.join(data_path, 'raw_images')
features_path = os.path.join(data_path, 'image_features')

print(data_path)

file_paths = []
for root, dirs, files in os.walk(image_path):
    for name in files:
        if name.endswith('JPEG'):
            file_path = os.path.join(root, name)
            file_paths.append(file_path)

# 图像编码时的最大分辨率，防止内存不够
max_res = 512

cache_id = 0
cache_datas = []
for id in tqdm(range(len(file_paths))):
    file_path = file_paths[id]
    image = Image.open(file_path)
    
    # resize too large image
    higher_res = max(image.size)
    if higher_res >= max_res:
        scale = higher_res / max_res
        image = image.resize((int(image.size[0] / scale), int(image.size[1] / scale)), Image.Resampling.NEAREST)

    image_feature = clip_encoder.encode_img(image)
    data = {
        'id' : id,
        'path' : file_path,
        'features' : image_feature
    }
    cache_datas.append(data)

    if len(cache_datas) > 1000:
        with open(os.path.join(features_path, f'cache_feature{cache_id}.bin'), 'wb') as fp:
            pickle.dump(cache_datas, fp)
            cache_id += 1
            cache_datas.clear()

with open(os.path.join(features_path, f'cache_feature{cache_id}.bin'), 'wb') as fp:
    pickle.dump(cache_datas, fp)